import json
import os


def data_convert_for_aishell1(aishell1_dir='/home/backup_nfs2/nfs1_data/data_aishell/transcript/'):

    for root, dirs, files in os.walk(aishell1_dir):
        total_res_dict = {}
        for dir_name in dirs:
            if dir_name.startswith('.'):
                continue
            json_path = f'./output/aishell_{dir_name}.json'
            res_dict = {}
            os.makedirs(os.path.dirname(json_path), exist_ok=True)
            text_file_path = os.path.join(root, dir_name, 'text')
            with open(text_file_path, 'r', encoding='utf-8') as f:
                lines = f.readlines()
            for line in lines:
                line = line.strip().split(" ")
                id = line[0]
                text_label = "".join(line[1:])
                res_dict[id] = {'transcription': text_label}
                total_res_dict[id] = {'transcription': text_label}
            wav_scp_path = os.path.join(root, dir_name, 'wav.scp')
            with open(wav_scp_path, 'r', encoding='utf-8') as f:
                lines = f.readlines()
            for line in lines:
                line = line.strip().split(" ")
                id = line[0]
                wav_path = line[1]
                res_dict[id]['wav'] = wav_path
                total_res_dict[id]['wav'] = wav_path
            with open(json_path, 'w', encoding='utf-8') as f:
                json.dump(res_dict, f, ensure_ascii=False, indent=4)
        with open('./output/aishell.json', 'w', encoding='utf-8') as f:
            json.dump(total_res_dict, f, ensure_ascii=False, indent=4)
        break


def make_fake_data_for_windows(small_aishell_dir='E:\gengxuelong_study\server_local_adapter\\ai\data\small_aishell'):
    for root, dirs, files in os.walk(small_aishell_dir):
        total_res_dict = {}
        for dir_name in dirs:
            print(dir_name)
            if dir_name.startswith('.'):
                continue
            json_path = f'./output/aishell_{dir_name}.json'
            res_dict = {}
            os.makedirs(os.path.dirname(json_path), exist_ok=True)
            wav_dir = os.path.join(root, dir_name, )
            for wav_name in os.listdir(wav_dir):
                wav_id = wav_name.split('.wav')[0]
                wav_path = os.path.join(wav_dir, wav_name)
                res_dict[wav_id] = {'transcription': "我是耿雪龙大家好哈哈哈你是谁", 'wav': wav_path}
                total_res_dict[wav_id] = {'transcription': "我是耿雪龙大家好哈哈哈你是谁", 'wav': wav_path}
            with open(json_path, 'w', encoding='utf-8') as f:
                json.dump(res_dict, f, ensure_ascii=False, indent=4)
            print(res_dict)
        print("total_res_dict", total_res_dict)
        with open('./output/aishell.json', 'w', encoding='utf-8') as f:
            json.dump(total_res_dict, f, ensure_ascii=False, indent=4)
        break


if __name__ == '__main__':
    data_convert_for_aishell1()
    # make_fake_data_for_windows()